import matplotlib.pyplot as plt
import numpy as np
from matplotlib.widgets import RectangleSelector    # 从小工具里面导入矩形选择器


class Mandbrot():
    def __init__(self, da):
        self.N = 5
        self.da = da
        self.isdrawing = False
        self.image = None
        self.CurCmap = "gist_rainbow"
        self.ax = plt.axes([0, 0, 1, 1], facecolor=da.bk)
        plt.axis("off")
        plt.axis("equal")
        plt.axis("tight")
        plt.axis("square")
        # 矩形选择器
        self.RS = RectangleSelector(self.ax, self.select_callback, 
                                    drawtype='box', useblit=True, 
                                    button=[1, 3], 
                                    minspanx=5, minspany=5, 
                                    rectprops=dict(facecolor='#5ff7e2', edgecolor='red', alpha=0.318, lw=3, fill=True), 
                                    spancoords='pixels', interactive=False)
        self.mdb(da.home[0], da.home[1], da.home[2], da.home[3], self.CurCmap)

        # self.mdb(da.g2[0], da.g2[1], da.g2[2], da.g2[3], self.CurCmap, 2)

    def mdb(self, xmin, xmax, ymin, ymax, cm):
        if self.isdrawing:
            return
        Nmax = 50
        print("{0} {1} {2} {3}".format(xmin, xmax, ymin, ymax))

        # self.ax.set_title("3434343443")
        self.ax.set_xlim([xmin, xmax])
        self.ax.set_ylim([ymin, ymax])
        # 离散化
        x0 = np.linspace(xmin, xmax, self.da.sw-2*self.da.margin)
        y0 = np.linspace(ymin, ymax, self.da.sw-2*self.da.margin)
        # 二维函数的离散化
        x, y = np.meshgrid(x0, y0)
        z = x + y*1j
        r = np.size(z, 0)   # rows
        c = np.size(z, 1)   # column
        w = np.zeros((r, c))
        k = w   # k是每个离散点的计数器
        N = 0
        while N < Nmax:
            w = w ** self.N + z  #
            N += 1
            t = (k==0) & (abs(w)>4)   # 逃逸的存在t当中
            k[t] = N    # w的模大于4的点， 并且该点的计数为0的话， 设计其计数为N
        k[k==0] = Nmax

        if self.image:  # 清除前面的图片
            del self.image
        self.image = plt.pcolormesh(x, y, k, edgecolor='none', cmap=cm)
        plt.savefig('temp.png')
        self.isdrawing = False
        plt.gcf().canvas.draw_idle()

    def select_callback(self, eclick, erelease):
        x1, y1 = eclick.xdata, eclick.ydata
        print("x1:{0}, y1:{1}".format(x1, y1))
        x2, y2 = erelease.xdata, erelease.ydata
        print("x2:{0}, y2:{1}".format(x2, y2))
        self.mdb(x1, x2, y1, y2, self.CurCmap)
        self.ax.figure.canvas.draw_idle()
